package com.wzy.oa_sys.config;

import com.wzy.oa_sys.service.ChatAssistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration(proxyBeanMethods = false)
public class LLMConfig {
    @Bean
    public ChatLanguageModel chatLanguageModel(){
        return OpenAiChatModel.builder()
                .apiKey("sk-H4bTxgol1hy7tSWwMQo1mh9EUOSJ2IjzSBkSvlJFFwwDBXsm")//模型api key
                .modelName("deepseek-v3")//模型名称
                .baseUrl("https://api.nuwaapi.com/v1")//模型url地址,提醒json问题则更换url地址
                .build();
    }

    /**
     * 嵌入存储 (简易内存存储)
     *
     * @return {@link InMemoryEmbeddingStore }<{@link TextSegment }>
     */
    @Bean
    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean
    public ChatAssistant assistant(ChatLanguageModel chatLanguageModel, EmbeddingStore<TextSegment> embeddingStore) {
        DefaultRetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .queryTransformer(new CompressingQueryTransformer(chatLanguageModel))  // 查询增强
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore)) // 内容源 单个直接配置
                .queryRouter(new DefaultQueryRouter(EmbeddingStoreContentRetriever.from(embeddingStore)))// 多个内容源，路由
                .contentInjector(new DefaultContentInjector())   // 结果提示词注入
                .build();

        return AiServices.builder(ChatAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
//                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .retrievalAugmentor(retrievalAugmentor)
                .build();
    }
}
